#!/usr/bin/env python
# -*- encoding: UTF-8 -*-
from __future__ import print_function
import argparse
import numpy as np
import cv2
import os
import glob

def calibration_from(w,h,path):
    # termination criteria 迭代终止标准
    criteria = (cv2.TERM_CRITERIA_EPS + cv2.TERM_CRITERIA_MAX_ITER, 30, 0.001)

    # prepare object points, like (0,0,0), (1,0,0), (2,0,0) ....,(6,5,0)
    # 准备目标点，三维空间中的，以棋盘格的为单位长度
    objp = np.zeros((h*w,3), np.float32)
    objp[:,:2] = np.mgrid[0:w,0:h].T.reshape(-1,2)

    # Arrays to store object points and image points from all the images.
    # 从所有图像中储存所有的目标点和图像点
    objpoints = [] # 3d point in real world space
    imgpoints = [] # 2d points in image plane.

    images = glob.glob(os.path.join(path,'*.jpg'))

    for fname in images:
        img = cv2.imread(fname)
        gray = cv2.cvtColor(img,cv2.COLOR_BGR2GRAY)

        # Find the chess board corners
        # 寻找标定角点
        ret, corners = cv2.findChessboardCorners(gray, (w,h),None)

        # If found, add object points, image points (after refining them)
        # 如果找到角点，添加目标点和图像点
        if ret == True:
            objpoints.append(objp)
            cv2.cornerSubPix(gray,corners,(11,11),(-1,-1),criteria) # 经试验发现 cornerSubPix 就地改变了 corners
            imgpoints.append(corners)
            # Draw and display the corners
            # 绘制并显示角点
            #img = cv2.drawChessboardCorners(img, (w,h), corners,ret)
            cv2.drawChessboardCorners(img, (w,h), corners,ret) # 经试验发现 drawChessboardCorners 就地改变了img
            cv2.imshow('image',img)
            cv2.waitKey(100)

    return cv2.calibrateCamera(objpoints, imgpoints, gray.shape[::-1],None,None)

def calibration_to(w,h,source,path):
    cap = cv2.VideoCapture(source)
    criteria = (cv2.TERM_CRITERIA_EPS + cv2.TERM_CRITERIA_MAX_ITER, 30, 0.001)
    num = 0
    while True:
        ret, frame = cap.read()
        if not ret:
            break
        show_img = frame.copy()
        gray = cv2.cvtColor(show_img,cv2.COLOR_BGR2GRAY)
        ret, corners = cv2.findChessboardCorners(gray, (w,h),None)
        if ret:
            cv2.cornerSubPix(gray,corners,(11,11),(-1,-1),criteria) # 经试验发现 cornerSubPix 就地改变了 corners
            cv2.drawChessboardCorners(show_img, (w,h), corners,ret) # 经试验发现 drawChessboardCorners 就地改变了img

        cv2.imshow('image', show_img)
        ch = cv2.waitKey(1) & 0xFF
        if ch == ord(' '):
            cv2.imwrite(os.path.join(path,"img_%d.jpg") % num,frame)
            num += 1

        if ch == 27: # Esc
            break

    return calibration_from(w,h,path)

def show(w,h,mtx, dist, source = 0):
    cap = cv2.VideoCapture(source)
    criteria = (cv2.TERM_CRITERIA_EPS + cv2.TERM_CRITERIA_MAX_ITER, 30, 0.001)

    objp = np.zeros((h*w,3), np.float32)
    objp[:,:2] = np.mgrid[0:w,0:h].T.reshape(-1,2)
    while True:
        ret, frame = cap.read()
        if not ret:
            break
        show_img = frame.copy()
        gray = cv2.cvtColor(show_img,cv2.COLOR_BGR2GRAY)
        ret, corners = cv2.findChessboardCorners(gray, (w,h),None)
        if ret:
            cv2.cornerSubPix(gray,corners,(11,11),(-1,-1),criteria) # 经试验发现 cornerSubPix 就地改变了 corners
            cv2.drawChessboardCorners(show_img, (w,h), corners,ret) # 经试验发现 drawChessboardCorners 就地改变了img
            objpoints = objp
            imgpoints = corners
            ret,rvec,tvec = cv2.solvePnP(objpoints, imgpoints,mtx,dist)
            if ret:
                print("rvec:",rvec)
                print("tvec:",tvec)
        cv2.imshow('image', show_img)
        ch = cv2.waitKey(1) & 0xFF
        if ch == 27:
            break

def main(w,h,source,path):
    if source == 'jpg':
        ret, mtx, dist, rvecs, tvecs = calibration_from(w,h,path)
    else:
        ret, mtx, dist, rvecs, tvecs = calibration_to(w,h,int(source),path)

    print(u"相机矩阵:")
    print(mtx)
    print(u"畸变系数:")
    print(dist)
    print(u"各系数值:")
    fx,fy,cx,cy = mtx[0,0],mtx[1,1],mtx[0,2],mtx[1,2]
    k1,k2,p1,p2,k3 = dist[0]
    K = {
        "fx": fx,
        "fy": fy,
        "cx": cx,
        "cy": cy,
        "k1": k1,
        "k2": k2,
        "p1": p1,
        "p2": p2,
        "k3": k3
    }
    for k in K:
        print(k,'=',K[k])
    if source != "jpg":
        show(w,h,mtx, dist,int(source))
    else:
        show(w,h,mtx, dist,0)

if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument("--w", type=int, default=12,
                        help="ChessBoard Width")
    parser.add_argument("--h", type=int, default=9,
                        help=u"ChessBoard Height")
    parser.add_argument("--source", type=str, default="jpg",
                        help=u"jpg or cam index")
    parser.add_argument("--path", type=str, default=os.curdir,
                        help="ChessBoard Images Path")
    args = parser.parse_args()
    main(args.w, args.h, args.source, args.path)
    cv2.destroyAllWindows()